#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
列车速度统计分析脚本
根据VINFO数组中的速度信息生成统计图表
"""

import argparse
import numpy as np
import matplotlib.pyplot as plt
from utils.h5data import read_h5_data
from utils.plot import *

def load_vinfo_data(h5_file):
    """
    从HDF5文件中加载VINFO数据
    
    Parameters:
    h5_file (str): HDF5文件路径
    
    Returns:
    tuple: 包含VINFO数组和其他相关信息的元组
    """
    # 读取VINFO数据
    VINFO = read_h5_data(h5_file, keys=['info/VINFO'], FILL_NONE=True)[0]
    
    if VINFO is None:
        raise ValueError("无法从文件中读取VINFO数据")
    
    return VINFO


def analyze_velocity_statistics(VINFO):
    """
    分析列车速度统计数据
    
    Parameters:
    VINFO (numpy.ndarray): VINFO数组，形状为(ne, 7)
    
    Returns:
    dict: 包含统计结果的字典
    """
    # VINFO结构说明:
    # VINFO[:, 0]: 方向标识 (0=无信号, 1=S2N, 2=N2S, 3=双向)
    # VINFO[:, 1:4]: 正向(南向北)速度信息 [v, a, e] 
    # VINFO[:, 4:7]: 负向(北向南)速度信息 [v, a, e]
    
    # 提取有效的速度数据
    # 正向速度(S2N)
    valid_s2n_mask = (VINFO[:, 0] == 1) | (VINFO[:, 0] == 3)  # S2N或双向
    s2n_velocities = VINFO[valid_s2n_mask, 1]  # 正向速度值
    s2n_velocities = s2n_velocities[s2n_velocities != 0]  # 去除无效值
    
    # 负向速度(N2S)
    valid_n2s_mask = (VINFO[:, 0] == 2) | (VINFO[:, 0] == 3)  # N2S或双向
    n2s_velocities = VINFO[valid_n2s_mask, 4]  # 负向速度值
    n2s_velocities = n2s_velocities[n2s_velocities != 0]  # 去除无效值
    n2s_velocities = np.abs(n2s_velocities)  # 取绝对值以便于比较
    
    # 方向统计
    no_signal_count = np.sum(VINFO[:, 0] == 0)      # 无信号
    s2n_only_count = np.sum(VINFO[:, 0] == 1)       # 仅S2N
    n2s_only_count = np.sum(VINFO[:, 0] == 2)       # 仅N2S
    both_direction_count = np.sum(VINFO[:, 0] == 3) # 双向
    
    return {
        's2n_velocities': s2n_velocities,
        'n2s_velocities': n2s_velocities,
        'direction_stats': {
            'no_signal': no_signal_count,
            's2n_only': s2n_only_count,
            'n2s_only': n2s_only_count,
            'both_direction': both_direction_count
        }
    }


def plot_velocity_statistics(stats_data, output_file=None):
    """
    绘制速度统计图
    
    Parameters:
    stats_data (dict): analyze_velocity_statistics函数返回的统计结果
    output_file (str): 输出图像文件路径，如果为None则显示图形
    """
    fig, axes = plt.subplots(2, 2, figsize=(6,5))
    fig.suptitle('列车速度统计分析', fontsize=16)
    
    # 1. 正向(S2N)速度直方图
    ax1 = axes[0, 0]
    if len(stats_data['s2n_velocities']) > 0:
        ax1.hist(stats_data['s2n_velocities'], bins=50, alpha=0.7, color='blue', edgecolor='black')
        ax1.axvline(np.mean(stats_data['s2n_velocities']), color='red', linestyle='--', 
                   label=f"平均速度: {np.mean(stats_data['s2n_velocities']):.2f} m/s")
        ax1.set_xlabel('速度 (m/s)')
        ax1.set_ylabel('频次')
        ax1.set_title('正向(S2N)速度分布')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
    else:
        ax1.text(0.5, 0.5, '无正向速度数据', ha='center', va='center', transform=ax1.transAxes)
        ax1.set_title('正向(S2N)速度分布')
    
    # 2. 负向(N2S)速度直方图
    ax2 = axes[0, 1]
    if len(stats_data['n2s_velocities']) > 0:
        ax2.hist(stats_data['n2s_velocities'], bins=50, alpha=0.7, color='orange', edgecolor='black')
        ax2.axvline(np.mean(stats_data['n2s_velocities']), color='red', linestyle='--', 
                   label=f"平均速度: {np.mean(stats_data['n2s_velocities']):.2f} m/s")
        ax2.set_xlabel('速度 (m/s)')
        ax2.set_ylabel('频次')
        ax2.set_title('负向(N2S)速度分布 (绝对值)')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
    else:
        ax2.text(0.5, 0.5, '无负向速度数据', ha='center', va='center', transform=ax2.transAxes)
        ax2.set_title('负向(N2S)速度分布 (绝对值)')
    
    # 3. 双方向速度对比箱线图
    ax3 = axes[1, 0]
    data_to_plot = []
    labels = []
    if len(stats_data['s2n_velocities']) > 0:
        data_to_plot.append(stats_data['s2n_velocities'])
        labels.append('正向(S2N)')
    if len(stats_data['n2s_velocities']) > 0:
        data_to_plot.append(stats_data['n2s_velocities'])
        labels.append('负向(N2S)')
    
    if len(data_to_plot) > 0:
        box_plot = ax3.boxplot(data_to_plot, labels=labels, patch_artist=True)
        # 设置箱体颜色
        colors = ['lightblue', 'lightcoral']
        for patch, color in zip(box_plot['boxes'], colors[:len(box_plot['boxes'])]):
            patch.set_facecolor(color)
        ax3.set_ylabel('速度 (m/s)')
        ax3.set_title('双方向速度对比')
        ax3.grid(True, alpha=0.3)
    else:
        ax3.text(0.5, 0.5, '无速度数据', ha='center', va='center', transform=ax3.transAxes)
        ax3.set_title('双方向速度对比')
    
    # 4. 方向分布饼图
    ax4 = axes[1, 1]
    directions = ['无信号', '仅正向(S2N)', '仅负向(N2S)', '双向']
    counts = [
        stats_data['direction_stats']['no_signal'],
        stats_data['direction_stats']['s2n_only'],
        stats_data['direction_stats']['n2s_only'],
        stats_data['direction_stats']['both_direction']
    ]
    
    # 只显示计数大于0的部分
    filtered_labels = [label for label, count in zip(directions, counts) if count > 0]
    filtered_counts = [count for count in counts if count > 0]
    filtered_colors = ['gray', 'blue', 'orange', 'green']
    filtered_colors = [color for color, count in zip(filtered_colors, counts) if count > 0]
    
    if len(filtered_counts) > 0:
        wedges, texts, autotexts = ax4.pie(filtered_counts, labels=filtered_labels, autopct='%1.1f%%', 
                                          colors=filtered_colors, startangle=90)
        ax4.set_title('列车方向分布')
    else:
        ax4.text(0.5, 0.5, '无事件数据', ha='center', va='center', transform=ax4.transAxes)
        ax4.set_title('列车方向分布')
    
    plt.tight_layout()
    
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"速度统计图已保存至: {output_file}")
    else:
        plt.show()


def main():
    parser = argparse.ArgumentParser(description='列车速度统计分析')
    parser.add_argument('-input', required=True, help='输入的HDF5文件路径')
    parser.add_argument('-output', help='输出图像文件路径 (可选，默认显示图形)')
    
    args = parser.parse_args()
    
    try:
        # 加载数据
        print(f"正在加载数据文件: {args.input}")
        VINFO = load_vinfo_data(args.input)
        print(f"成功加载VINFO数据，形状: {VINFO.shape}")
        
        # 分析数据
        print("正在进行速度统计分析...")
        stats_data = analyze_velocity_statistics(VINFO)
        
        s2n_count = len(stats_data['s2n_velocities'])
        n2s_count = len(stats_data['n2s_velocities'])
        print(f"检测到正向(S2N)事件: {s2n_count} 个")
        print(f"检测到负向(N2S)事件: {n2s_count} 个")
        
        dir_stats = stats_data['direction_stats']
        print(f"方向统计:")
        print(f"  无信号: {dir_stats['no_signal']} 个")
        print(f"  仅正向(S2N): {dir_stats['s2n_only']} 个")
        print(f"  仅负向(N2S): {dir_stats['n2s_only']} 个")
        print(f"  双向: {dir_stats['both_direction']} 个")
        
        # 绘图
        print("正在生成统计图表...")
        plot_velocity_statistics(stats_data, args.output)
        
    except Exception as e:
        print(f"处理过程中出现错误: {e}")
        raise


if __name__ == '__main__':
    main()